#!/bin/bash

eth_name=$1
#获取cpu总数
total_cpus=$(grep -c processor /proc/cpuinfo)
#获取网卡对应的PCI设备ID
pci_id=$(ls -l /sys/class/net/ |grep $eth_name |awk -F/ '{print $(NF-3)}')
#获取PCI设备对应的msi_irqs
msi_irqs=$(ls "/sys/bus/pci/devices/$pci_id/msi_irqs")

function show_eth_irqs()
{
   eth_name_pci_id=`ls -l /sys/class/net/ |grep $eth_name |awk -F/ '{print $(NF-3)}'`
   echo "$eth_name pci bus is $eth_name_pci_id"
   echo "$eth_name irqs are: "
   irq_num=0
   for dir in `ls /sys/bus/pci/devices/$eth_name_pci_id/msi_irqs/ | sort -n`
   do
      if [ -d /proc/irq/$dir ];then
         cpu=`cat /proc/irq/$dir/smp_affinity_list`
         echo "irq number $dir is binded to $cpu"
         irq_num=$(($irq_num+1))
      fi
   done
   echo "total irqs number: $irq_num"
}

function set_irq_balance()
{
   eth_name_pci_id=`ls -l /sys/class/net/ |grep $eth_name |awk -F/ '{print $(NF-3)}'`
   irq_num=`cat /proc/interrupts | grep -E 'input|output'| awk -F":" '{print $1}' | wc -l`
   
   if [ $irq_num -ge $total_cpus ]; then
      #中断数大于等于CPU数时，1个CPU绑定1个或多个中断
      per_cpu_irqs=$(($irq_num / $total_cpus))
      core=0
      count=0
      for i in `cat /proc/interrupts | grep -E 'input|output'| awk -F":" '{print $1}'`
      do
         if echo "$msi_irqs" | grep -q "$i"; then
            echo "binding irq $i to cpu $core"
            echo $core > /proc/irq/$i/smp_affinity_list
            count=$(($count+1))
            if [ $(($count % $per_cpu_irqs)) -eq 0 ];then
               core=$(($core+1))
            fi

            if [ $core -gt $(($total_cpus-1)) ];then
               core=0
            fi
         fi
      done
   else
      #中断数小于CPU数时，多个CPU绑定1个中断
      first_irq=`cat /proc/interrupts | grep -E 'input|output'| awk -F":" '{print $1}' | sort -n | head -1`
      last_irq=`cat /proc/interrupts | grep -E 'input|output'| awk -F":" '{print $1}' | sort -n | tail -1`
      irq=$first_irq
      for core in `seq 0 1 $(($total_cpus-1))`
      do
         if echo "$msi_irqs" | grep -q "$i"; then
            echo "binding irq $irq to cpu $core"
            echo $core > /proc/irq/$irq/smp_affinity_list
         fi

            irq=$(($irq+1))
            if [ $irq -gt $last_irq ];then
               irq=$first_irq
            fi
      done
   fi
}

function set_irq_affinity_node()
{
   if [ "$1" == "numa0" ];then
      start_cpu=0
      end_cpu=$(($total_cpus/2-1))
   else
      start_cpu=$(($total_cpus/2))
      end_cpu=$(($total_cpus-1))
   fi

   core=$start_cpu
   for i in `cat /proc/interrupts | grep -E 'input|output'| awk -F":" '{print $1}'`
   do
      if echo "$msi_irqs" | grep -q "$i"; then
         echo "binding irq $i to cpu $core"
         echo $core > /proc/irq/$i/smp_affinity_list
         core=$(($core+1))
         if [ $core -gt $end_cpu ];then
            core=$start_cpu
         fi
       fi
   done
}

function set_irq_affinity_to_speccpu()
{
   CPURANGE=$1
   STARTCPU=`echo $CPURANGE | cut -d '-' -f 1`
   ENDCPU=`echo $CPURANGE | cut -d '-' -f 2`
   core=$STARTCPU

   for i in `cat /proc/interrupts | grep -E 'input'| awk -F":" '{print $1}'`
   do
      if echo "$msi_irqs" | grep -q "$i"; then
         echo "binding irq $i to cpu $core"
         echo $core > /proc/irq/$i/smp_affinity_list
         core=$((core+1))
         if [ $core -gt $ENDCPU ];then
             core=$STARTCPU
         fi
       fi
   done

   for i in `cat /proc/interrupts | grep -E 'output'| awk -F":" '{print $1}'`
   do
      echo "binding irq $i to cpu $core"
      echo $core > /proc/irq/$i/smp_affinity_list
      core=$((core+1))
      if [ $core -gt $ENDCPU ];then
         core=$STARTCPU
      fi
   done   
}

function usage()
{
   echo "Usage: 
      sh $0 eth_name [ show | balance | numa0 | numa1 | speccpu ]"
}

if [[ "$#" -ne 2 &&  "$#" -ne 3 ]]; then
  usage
  exit 1
fi

case $2 in
  show)
    echo "Executing command show irq"
    show_eth_irqs
    ;;
  balance)
    echo "Executing command set irq balance"
    set_irq_balance
    ;;
  numa0)
    echo "Executing command set irq numa0"
    set_irq_affinity_node numa0
    ;;
  numa1)
    echo "Executing command set irq numa1"
    set_irq_affinity_node numa1
    ;;
  speccpu)
    echo "Executing command set irq speccpu"
    set_irq_affinity_to_speccpu  $3
    ;;
  *)
    echo "Invalid command: $2"
    usage
    exit 1
    ;;
esac